#define _WIN32_WINNT 0x0600
#include <windows.h>
#include <winbase.h>
#include <tchar.h>

#include <cstdio>
#include <ctime>
#include <cstring>

bool
LoggedSetLockPagesPrivilege(HANDLE process);

int
main(int argc, char **args)
{
  // 引数チェック
  if (argc < 4) {
    _tprintf(_T("usage : %s (allocate size[Mbyte]) (sleep time[sec]) (mapped[true/false])\n"), args[0]);
    return 1;
  }

  int memoryRequest = atoi(args[1]) * 1024 * 1024;
  int sleepSecs = atoi(args[2]) * 1000;
  bool mapped = false;
  if (stricmp("true", args[3]) == 0)
  {
    mapped = true;
  }

  _tprintf(_T("AWE allocate test\n  memory size : %d byte\n  sleep time : %d msec\n"), memoryRequest, sleepSecs);
  if (mapped)
  {
    _tprintf(_T("  mapped : true\n"));
  }
  else
  {
    _tprintf(_T("  mapped : false\n"));
  }

  BOOL result;

  // OSで統一されているページサイズの取得
  SYSTEM_INFO sysinfo;
  GetSystemInfo(&sysinfo);
  _tprintf(_T("This computer has page size %lu.\n"), sysinfo.dwPageSize);

  // 要求されているメモリサイズに該当する物理ページ数を計算する
  ULONG_PTR numberOfPages = memoryRequest / sysinfo.dwPageSize;
  _tprintf(_T("Requesting %lu pages of memory.\n"), numberOfPages);

  // PNF配列のデータサイズを算出し、確保する
  int PFNArraySize = numberOfPages * sizeof (ULONG_PTR);
  _tprintf(_T("Requesting a PFN array of %d bytes.\n"), PFNArraySize);

  ULONG_PTR *pfns = (ULONG_PTR *)HeapAlloc(GetProcessHeap(), 0, PFNArraySize);
  if (pfns == NULL)
  {
    _tprintf(_T("Failed to allocate on heap.\n"));
    return 1;
  }

  HANDLE process = GetCurrentProcess();

  // AWE用の権限を取得する
  if ( ! LoggedSetLockPagesPrivilege(process))
  {
    return 1;
  }

  // AWEを使い物理メモリを確保する
  {
    ULONG_PTR numberOfPagesInitial = numberOfPages;

    result = AllocateUserPhysicalPages(process, &numberOfPages, pfns);
    if (result != TRUE)
    {
      _tprintf(_T("Cannot allocate physical pages (%lu)\n"), GetLastError());
      return 1;
    }

    if (numberOfPagesInitial != numberOfPages)
    {
      _tprintf(_T("Allocated only %lu pages.\n"), numberOfPages);
      return 1;
    }
  }

  // 仮想メモリを確保する
  PVOID memReserved;            // AWE window
  memReserved = VirtualAlloc(NULL, memoryRequest, MEM_RESERVE | MEM_PHYSICAL, PAGE_READWRITE);
  if (memReserved == NULL)
  {
    _tprintf(_T("Cannot reserve memory.\n"));
    return 1;
  }

  if (! mapped)
  {
    Sleep(sleepSecs);
  }
  else
  {
    // AWEで確保した物理ページをプロセス空間にマップする
    result = MapUserPhysicalPages(memReserved, numberOfPages, pfns);
    if (result != TRUE)
    {
      _tprintf(_T("MapUserPhysicalPages failed (%lu)\n"), GetLastError());
      return 1;
    }

    Sleep(sleepSecs);

    // AWEで確保した物理ページをプロセス空間から切り離す
    result = MapUserPhysicalPages(memReserved, numberOfPages, NULL);
    if (result != TRUE)
    {
      _tprintf(_T("MapUserPhysicalPages failed (%lu)\n"), GetLastError());
      return 1;
    }
  }

  // AWEで確保した物理ページを開放する
  result = FreeUserPhysicalPages(process, &numberOfPages, pfns);
  if (result != TRUE)
  {
    _tprintf(_T("Cannot free physical pages, error %lu.\n"), GetLastError());
    return 1;
  }

  // 仮想メモリを開放する
  result = VirtualFree(memReserved, 0, MEM_RELEASE);

  // PFN用の領域を開放する
  result = HeapFree(GetProcessHeap(), 0, pfns);
  if (result != TRUE)
  {
    _tprintf(_T("Call to HeapFree has failed (%lu)\n"), GetLastError());
  }

  return 0;
}

bool
LoggedSetLockPagesPrivilege(HANDLE process)
{
  struct {
    DWORD count;
    LUID_AND_ATTRIBUTES privilege[1];
  } info;

  BOOL result;

  // プロセスの権限トークンを取得する
  HANDLE token;
  result = OpenProcessToken(process, TOKEN_ADJUST_PRIVILEGES, &token);
  if (result != TRUE)
  {
    _tprintf(_T("Cannot open process token.\n"));
    return false;
  }

  info.count = 1;
  info.privilege[0].Attributes = SE_PRIVILEGE_ENABLED;

  // メモリ関連のLUIDを取得する
  result = LookupPrivilegeValue(NULL, SE_LOCK_MEMORY_NAME, &(info.privilege[0].Luid));
  if (result != TRUE)
  {
    _tprintf(_T("Cannot get privilege for %s.\n"), SE_LOCK_MEMORY_NAME);
    return false;
  }

  // メモリ関連の権限を調整する
  result = AdjustTokenPrivileges(token, FALSE, (PTOKEN_PRIVILEGES)&info, 0, NULL, NULL);
  if (result != TRUE)
  {
    _tprintf(_T("Cannot adjust token privileges (%lu)\n"), GetLastError());
    return false;
  }

  if (GetLastError() != ERROR_SUCCESS)
  {
    _tprintf(_T("Cannot enable the SE_LOCK_MEMORY_NAME privilege; "));
    _tprintf(_T("please check the local policy.\n"));
    return false;
  }

  CloseHandle(token);

  return true;
}
